//
// Created by vivi on 27/07/2017.
//

#include "HairCore.h"

namespace XRwy {
	using namespace std;
	using namespace Eigen;

    vector<Vector3f> filter_kernel(const vector<Vector3f> &values, float alpha) {
        const float beta = min<float>(1.0f, 1.0f - exp(-1.0f / alpha));
        const float beta2 = beta * beta;
        const float _1_minus_beta = 1 - beta;
        const float _1_minus_beta_2 = _1_minus_beta * _1_minus_beta;

        Vector3f lst, lst2;
        lst2 = lst = (values.back() - values.front()) / (values.size() - 1);

        vector<Vector3f> ret;
        ret.reserve(values.size());
        ret.push_back(values[0]);
        for (size_t i = 1; i < values.size(); ++i) {
            Vector3f d = 2 * _1_minus_beta * lst - _1_minus_beta_2 * lst2 + beta2 * (values[i] - values[i - 1]);
            lst2 = lst;
            lst = d;

            ret.push_back(ret[i - 1] + d);
        }

        return std::move(ret);
    }

    struct DFTLStictionLinkHash {
        size_t operator()(const HairDFTLStictionLinkInfo &x) {
            return x.particles[0].strand_idx ^ x.particles[0].rel_idx ^ x.particles[1].strand_idx ^ x.particles[1].rel_idx;
        }
    };
}

namespace XRwy {
	using namespace std;
	using namespace Eigen;

    void HairDFTLStrand::get_frames() {
        frames.resize(particle_count);

        frames[0] = Matrix3f::Zero();
        for (size_t i = 1; i < filtered_pos_array.size(); ++i) {
            Vector3f d = filtered_pos_array[i] - filtered_pos_array[i - 1];

            Vector3f aligned_x = d;
            Vector3f aligned_y = aligned_x.cross(Vector3f::UnitX());
            Vector3f aligned_z = aligned_x.cross(aligned_y);

            aligned_x.normalize();
            aligned_y.normalize();
            aligned_z.normalize();

            frames[i] << aligned_x, aligned_y, aligned_z;
        }
    }

    void HairDFTLStrand::get_rel_coords() {
        rel_coords.resize(particle_count);

        rel_coords[0] = Vector3f::Zero();
        for (size_t i = 1; i < pos_array.size(); ++i)
            rel_coords[i] = frames[i].transpose() * (pos_array[i] - filtered_pos_array[i]);
    }

    void HairDFTLStrand::get_pos_array() {
        for (size_t i = 0; i < particle_count; ++i)
            pos_array[i] = filtered_pos_array[i] + frames[i] * rel_coords[i];
    }

    void HairDFTLStrand::get_filtered_pos() {
        for (size_t i = 0; i < particle_count; ++i)
            filtered_pos_array[i] = pos_array[i] - frames[i] * rel_coords[i];
    }

    HairDFTLSimulator::HairDFTLSimulator(const HairReader &hair_reader, const HairDFTLConfiguration &conf_):
            conf(conf_), particle_count((size_t)hair_reader.get_particle_count()), strand_count((size_t)hair_reader.get_strand_count()) {

        strands.reserve((size_t)hair_reader.get_strand_count());

        size_t cur = 0;
        for (const auto & strand_struct : hair_reader.get_strands()) {
            strands.emplace_back();
            auto & strand = strands.back();

            strand.particle_begin_index = cur;
            strand.particle_count = strand_struct.size();

            strand.pos_array = strand_struct;
            strand.vel_array = vector<Vector3f>(strand_struct.size(), Vector3f::Zero());
            strand.filtered_pos_array = filter_kernel(strand.pos_array, conf.filter_alpha);
            strand.rest_filtered_pos_array = strand.filtered_pos_array;

            strand.get_rest_length();
            strand.get_frames();
            strand.get_rel_coords();

            cur += strand.particle_count;
        }

        //for stiction
        aligned_x_particle_indexes.reserve((size_t)hair_reader.get_particle_count());
        for (size_t i = 0; i < strands.size(); ++i)
            for (size_t j = 0; j < strands[i].particle_count; ++j)
                aligned_x_particle_indexes.push_back(HairDFTLParticleInfo{i, j});

        stiction_springs.clear();
    }

    void HairDFTLStrand::get_rest_length() {
        rest_length.resize(particle_count - 1);
        for (size_t i = 0; i < particle_count - 1; ++i)
            rest_length[i] = (filtered_pos_array[i + 1] - filtered_pos_array[i]).norm();
    }

    void HairDFTLSimulator::on_frame(const HairSimulatorData &data){
        Vector3f gravity;
        gravity << 0.0f, -9.8f, 0.0f;

        data.sdf_ptr->update(data.sdf_world2local, data.delta_time);

//        for (auto & strand : strands) {
//
//            vector<Vector3f> pos1 = strand.filtered_pos_array;
//            for (size_t i = 1; i < strand.filtered_pos_array.size(); ++i)
//                pos1[i] += data.delta_time * strand.vel_array[i] + (data.delta_time * data.delta_time) * gravity;
//            pos1[0] = data.motion_transform * strand.rest_filtered_pos_array[0];
//
//            vector<Vector3f> pos2 = pos1;
//            for (size_t i = 1; i < pos2.size(); ++i) {
//                pos2[i] = pos2[i - 1] + strand.rest_length[i - 1] * (pos2[i] - pos2[i - 1]).normalized();
//            }
//
//            for (size_t i = 0; i < strand.filtered_pos_array.size(); ++i)
//                strand.vel_array[i] = (1.0f / data.delta_time) * (pos2[i] - strand.filtered_pos_array[i] - conf.damping * (pos2[i] - pos1[i]));
//
//            strand.filtered_pos_array = pos2;
//        }

        //resolve_time_intergration(data);

        const float t = data.delta_time / conf.npass;
        const float npass_1 = 1.0f / static_cast<float>(conf.npass);

        //pre compute the difference position for each strand root
        vector<Vector3f> dpos(strands.size());
        for (size_t i = 0; i < strands.size(); ++i)
            dpos[i] = npass_1 * (data.motion_transform * strands[i].rest_filtered_pos_array[0] - strands[i].filtered_pos_array[0]);

        for (size_t _ = 0; _ < conf.npass; ++_) {

            tbb::parallel_for(size_t(0), strands.size(), [this, t, &dpos] (size_t i) {
                resolve_time_intergration_for_strand(strands[i], strands[i].filtered_pos_array[0] + dpos[i], t);
            });

            if (conf.enable_hair_stiction)
                resolve_stiction(data);

            if (conf.enable_body_collision)
                resolve_collision(data);
        }

        get_pos_from_filtered();
    }

    HairDFTLSimulator::~HairDFTLSimulator() {
    }

    void HairDFTLSimulator::resolve_time_intergration(const HairSimulatorData &data) {

        const float t = data.delta_time / conf.npass;

        tbb::parallel_for(size_t(0), strands.size(), [this, t, &data](size_t i){
            auto & strand = strands[i];

            Vector3f new_root_pos = data.motion_transform * strand.rest_filtered_pos_array[0];
            Vector3f current_root_pos = strand.filtered_pos_array[0];
            Vector3f dpos = (new_root_pos - current_root_pos) / conf.npass;

            for (size_t i = 0; i < conf.npass; ++i) {
                current_root_pos += dpos;
                resolve_time_intergration_for_strand(strand, current_root_pos, t);
            }
        });
    }

    void HairDFTLSimulator::resolve_time_intergration_for_strand(HairDFTLStrand &strand, const Vector3f &strand_root_pos, const float t) {
        size_t dim = strand.particle_count * 3;

        SparseMatrix<float> A(dim, dim);
        VectorXf b = VectorXf::Zero(dim);

        Triplet<float> *A_triplets = new Triplet<float>[(strand.particle_count - 1) * 36 + dim];
        size_t A_triplets_count = 0;

        //add velocity to b
        for (size_t i = 0, bi = 0; i < strand.particle_count; ++i, bi += 3)
            b.segment<3>(bi) += strand.vel_array[i] + t * conf.ext_acc;

        //add acceleration to b and d matrix to A
        for (size_t i = 1; i < strand.particle_count; ++i) {
            Vector3f vec = strand.filtered_pos_array[i] - strand.filtered_pos_array[i - 1];
            float l = vec.norm();
            Vector3f d = vec / l;
            float l0 = strand.rest_length[i - 1]; //the rest length of particle[i] and particle[i-1] is stored in rest_length[i - 1]

            //add acceleration to b
            Vector3f acc = (t * conf.k * (l / l0 - 1)) * d;
            b.segment<3>(i * 3 - 3) += acc;
            b.segment<3>(i * 3) -= acc;

            //add d matrix to A
            float factor = conf.k / l0 * t * t;
            float dm[9] = {
                    factor * d(0) * d(0), factor * d(0) * d(1), factor * d(0) * d(2),
                    factor * d(1) * d(0), factor * d(1) * d(1), factor * d(1) * d(2),
                    factor * d(2) * d(0), factor * d(2) * d(1), factor * d(2) * d(2)
            };

            size_t apply_positions[2] = { 3 * i, 3 * i - 3};
            for (size_t y : apply_positions)
                for (size_t x : apply_positions)
                    for (size_t dy = 0; dy < 3; ++dy)
                        for (size_t dx = 0; dx < 3; ++dx)
                            A_triplets[A_triplets_count++] = Triplet<float>(y + dy, x + dx, (y == x ? 1.0f : -1.0f) * dm[dy * 3 + dx]);
        }

        //add damping and identity matrix to A
        for (size_t i = 0; i < dim; ++i)
            A_triplets[A_triplets_count++] = Triplet<float>(i, i, 1 + conf.damping * t);

        A.setFromTriplets(A_triplets, A_triplets + A_triplets_count);

        ConjugateGradient<SparseMatrix<float>, Lower | Upper> cg;

        cg.compute(A);
        VectorXf new_vel = cg.solve(b);

        strand.vel_array[0] = (strand_root_pos - strand.filtered_pos_array[0]) / t;
        strand.filtered_pos_array[0] = strand_root_pos;
        for (size_t i = 1, vi = 3; i < strand.particle_count; ++i, vi += 3) {
            strand.vel_array[i] = new_vel.segment<3>(vi);
            strand.filtered_pos_array[i] += new_vel.segment<3>(vi) * t;
        }

        delete [] A_triplets;
    }

    void HairDFTLSimulator::resolve_collision(const HairSimulatorData &data) {
        for (auto & strand : strands) {
            for (size_t i = 1; i < strand.filtered_pos_array.size(); ++i)
                data.sdf_ptr->fix(strand.filtered_pos_array[i], strand.vel_array[i], data.delta_time);
        }
    }

    void HairDFTLSimulator::get_pos_from_filtered() {
        for (auto & strand : strands) {
            strand.get_frames();
            strand.get_pos_array();
        }
    }

    void HairDFTLSimulator::resolve_stiction(const HairSimulatorData &data) {
        //Remove some spring that has the distance larger than the invalid dr
        vector<HairDFTLStictionLinkInfo> new_stiction_springs;

        vector<vector<int>> linked_counts(strand_count);
        for (size_t i = 0; i < strand_count; ++i)
            linked_counts[i] = vector<int>(strands[i].particle_count, 0);

        unordered_set<HairDFTLStictionLinkInfo, DFTLStictionLinkHash> links;
        new_stiction_springs.reserve(particle_count * conf.max_replusion_link);

        for (const auto &spring : stiction_springs)
            if ((get_particle_filtered_pos(spring.particles[0]) - get_particle_filtered_pos(spring.particles[1])).norm() < conf.invalid_dr) {
                new_stiction_springs.push_back(spring);
                ++linked_counts[spring.particles[0].strand_idx][spring.particles[0].rel_idx];
                ++linked_counts[spring.particles[1].strand_idx][spring.particles[1].rel_idx];
                links.insert(spring);
            }

        stiction_springs = std::move(new_stiction_springs);

        //update aligned hireachy
        std::sort(aligned_x_particle_indexes.begin(), aligned_x_particle_indexes.end(), [this](const HairDFTLParticleInfo &lhs, const HairDFTLParticleInfo &rhs) {
            const Vector3f & pl = this->get_particle_filtered_pos(lhs);
            const Vector3f & pr = this->get_particle_filtered_pos(rhs);
            return pl(0) < pr(0);
        });

        for (int i = 0; i < this->particle_count; ++i) {
            HairDFTLParticleInfo index_i = aligned_x_particle_indexes[i];
            const Vector3f &pos_i = get_particle_filtered_pos(index_i);

            for (int j = i + 1; linked_counts[index_i.strand_idx][index_i.rel_idx] < conf.max_replusion_link && j < this->particle_count; ++j) {
                HairDFTLParticleInfo index_j = aligned_x_particle_indexes[j];
                const Vector3f &pos_j = get_particle_filtered_pos(index_j);

                if (pos_j(0) - pos_i(0) > conf.dr)
                    break;

                float dist = (pos_i - pos_j).norm();
                HairDFTLStictionLinkInfo link(index_i, index_j);

                if (dist < conf.dr
                    && index_i.strand_idx != index_j.strand_idx
                    && linked_counts[index_j.strand_idx][index_j.rel_idx] < conf.max_replusion_link
                    && links.find(link) == links.end()) {
                    ++linked_counts[index_i.strand_idx][index_i.rel_idx];
                    ++linked_counts[index_j.strand_idx][index_j.rel_idx];
                    //0 stands for no strand index
                    stiction_springs.push_back(link);
                }
            }
        }

        cout << "Total self-replusion link count is: " << stiction_springs.size() / 2 << endl;
        cout << "Average link count is: " << (double)(stiction_springs.size() / 2) / this->particle_count << endl;

        for (const auto & link : stiction_springs) {
            Vector3f & pos1 = get_particle_filtered_pos(link.particles[0]);
            Vector3f & pos2 = get_particle_filtered_pos(link.particles[1]);

            Vector3f d = pos2 - pos1;
            Vector3f impulse = (data.delta_time * conf.kr * (d.norm() / conf.dr - 1.0f)) * d.normalized();

            get_particle_vel(link.particles[0]) += impulse;
            get_particle_vel(link.particles[1]) -= impulse;
        }
    }
}